package com.wyt.common.file.video.h264Tortmp;

import org.apache.commons.lang3.StringUtils;
import org.bytedeco.javacpp.avcodec;
import org.bytedeco.javacv.FFmpegFrameGrabber;
import org.bytedeco.javacv.FFmpegFrameRecorder;
import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.FrameRecorder;
import org.java_websocket.client.WebSocketClient;
import org.java_websocket.handshake.ServerHandshake;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.PipedInputStream;
import java.io.PipedOutputStream;
import java.net.URI;
import java.nio.ByteBuffer;

public class H264PushRtmp extends Thread {
    private static Logger logger = LoggerFactory.getLogger(H264PushRtmp.class);

    private final PipedInputStream pis = new PipedInputStream();
    private final PipedOutputStream pos = new PipedOutputStream();
    private WebSocketClient webSocketClient = null;
    private final String deviceId;
    private volatile boolean running = true;
    private boolean flag = false;
    private String rtmpBaseAddress = "rtmp://nginx-flv.cube.svc.cluster.local:1935/live/";
    private String wsBaseAddress = "ws://dataaccess.dky.svc.cluster.local:8098/video/";

    public H264PushRtmp(String deviceId, String rtmpUrl, String wsVideoUrl) {
        this.deviceId = deviceId;
        if (StringUtils.isNotBlank(rtmpUrl)) this.rtmpBaseAddress = rtmpUrl;
        if (StringUtils.isNotBlank(wsVideoUrl)) this.wsBaseAddress = wsVideoUrl;
        try {
            // 使用管道流实现线程间通信
            pos.connect(pis);
        } catch (IOException e) {
            logger.error("pos connect pis error", e);
        }
    }

    /**
     * 将视频流写入管道
     */
    public void onMediaStream(byte[] data) {
        try {
            if (!flag) {
                logger.info("receive data...");
                flag = true;
            }
            pos.write(data);
        } catch (IOException e) {
            logger.error("write video data error.", e);
//            try {
//                pos.close();
//            } catch (IOException ex) {
//                logger.error("pos close error.{}", ex.getMessage());
//            }
        }
    }

    /**
     * 转流器, 指定format
     */
    public void recordPushWithFormat() {
        long startTime = 0;
        long videoTS;

        try {
            logger.info("grabber start ... {} ", deviceId);
            // 从管道流中读取视频流
            // maximumSize 设置为0，不设置会阻塞
//            FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(pis, 0);
            FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(pis);

            //此配置以减少启动时间！若不设置，启动最起码半分钟；
            //类似一个缓冲区，用来检测流的正确性，值越小，启动越快
            grabber.setOption("probesize", "1024*20");
            grabber.setFrameRate(30);
            grabber.setVideoBitrate(2000000);
            //阻塞式，直到通道有数据
            grabber.start();

            logger.info("grabber start suc. and start recorder ... {} ", deviceId);

            // 获取推流地址对应的recorder，用于录制视频
            FrameRecorder rtmpRecorder = getRecorder(rtmpBaseAddress + deviceId, "flv");

            rtmpRecorder.start();

            Frame grabframe;

            // 从视频流中捕获帧以录制视频
            while (running && (grabframe = grabber.grab()) != null) {

                if (startTime == 0) {
                    startTime = System.currentTimeMillis();
                }

                videoTS = 1000 * (System.currentTimeMillis() - startTime);

                // 推流到rtmp server
                if (videoTS > rtmpRecorder.getTimestamp()) {
                    rtmpRecorder.setTimestamp(videoTS);
                }
                rtmpRecorder.record(grabframe);

            }
        } catch (Exception e) {
            logger.error("record push error.", e);
        } finally {
            try {
                pis.close();
            } catch (IOException e) {
                logger.error("pis close error.{}", e.getMessage());
            }
        }

    }

    private FrameRecorder getRecorder(String address, String format) {
        FrameRecorder recorder = new FFmpegFrameRecorder(address, 720, 480);
        recorder.setInterleaved(true);
        recorder.setVideoOption("tune", "zerolatency");
        recorder.setVideoOption("preset", "ultrafast");
        recorder.setVideoOption("crf", "28");
        recorder.setVideoBitrate(2000000);//码率 越大越清晰
        recorder.setVideoCodec(avcodec.AV_CODEC_ID_H264);
        recorder.setFormat(format);//flv rtsp
        recorder.setFrameRate(30);//帧率 30fps 每秒包含的帧数 24-30越大越流畅
        recorder.setGopSize(60);//30*2 每60帧存在一个关键帧
        return recorder;
    }

    public void stopPush() {
        this.running = false;
        if (webSocketClient != null) webSocketClient.close();
        this.interrupt();
    }


    public WebSocketClient getWsClient() throws Exception {
        webSocketClient = new WebSocketClient(new URI(wsBaseAddress+deviceId)) {
            @Override
            public void onOpen(ServerHandshake serverHandshake) {
                System.out.println("建立链接");
            }

            @Override
            public void onMessage(String s) {
                System.out.println("接收消息：" + s);
            }

            @Override
            public void onMessage(ByteBuffer bytes) {
                System.out.println("byte len:" + bytes.array().length);
                try {
                    //将接收到的字节推送到管道
                    onMediaStream(bytes.array());
                } catch (Exception e) {
                    logger.info("推送rtmp异常:\n", e);
                }
            }

            @Override
            public void onClose(int i, String s, boolean b) {
                System.out.println("关闭客户端");
            }

            @Override
            public void onError(Exception e) {
                logger.error("error:", e);
            }
        };
        webSocketClient.connect();
        return webSocketClient;
    }


    @Override
    public void run() {
        try {
            //启动流捕获程序，持续捕获外部传入的数据
            recordPushWithFormat();
        } catch (Exception e) {
            logger.error("创建websocket异常", e);
        }
    }

    public static void main(String[] args) {
        H264PushRtmp rtmp = new H264PushRtmp("kvb2060004", "rtmp://127.0.0.1:1935/live/", "ws://127.0.0.1:8098/video/");
        new Thread(rtmp).start();

        try {
            //启动websocket监听视频流,将接收到的字节通过管道传输到推送线程中
            WebSocketClient wsClient = rtmp.getWsClient();
            while (!wsClient.isOpen()) {
                System.out.println(1);
                Thread.sleep(1000);
            }
            System.out.println("链接成功.....");
        } catch (Exception e) {

        }
    }


}
